import argparse
import os
import shutil
from pathlib import Path
import numpy as np
import torch
import yaml
from tqdm import tqdm
from models import transfer
from utils.datasets import create_dataloader
from utils.tools import estimate, plot_confusion_matrix, make_confusion_matrix
from utils.torch_utils import cuda2cpu, select_device


def test(opt, device):
    out = Path('inference/output')
    if os.path.exists(out):
        shutil.rmtree(out)  # delete output folder
    os.makedirs(out)  # make new output folder

    with open(opt.data,encoding="utf-8") as f:
        data_dict = yaml.load(f, Loader=yaml.FullLoader)  # model dict

    nc = int(data_dict['nc'])
    model = transfer.Model(opt.model_name, nc=nc, pretrained=False).to(device)  # create
    model.load_state_dict(torch.load(opt.weights, map_location=device)["model"], strict=False)
    model.eval()
    model.to(device)

    test_dataloader, test_dataset = create_dataloader(data_dict[opt.task], opt, None, data_dict, isTrain=False)

    preds = []
    gt_labels = []
    for i, (inputs, labels, paths) in tqdm(enumerate(test_dataloader), total=len(test_dataloader)):
        inputs = inputs.to(device)

        outputs = model(inputs)
        prob, pred = torch.max(outputs, 1)

        pred_array = cuda2cpu(pred).tolist()
        preds.extend(pred_array)
        gt_labels.extend(labels)

        if opt.save_error_img:
            errorimgpaths = np.array(paths)[np.array(cuda2cpu(pred)) != np.array(labels)]
            if len(paths) == 1 and len(errorimgpaths) != 0:
                shutil.copy(errorimgpath, out)
            else:
                for errorimgpath in errorimgpaths:
                    shutil.copy(errorimgpath, out)

    accuracy, precision, recall, f1score, cm = estimate(gt_labels, preds)
    plot_confusion_matrix(cm, labels=data_dict["names"], cm_path=out)
    make_confusion_matrix(cm, cm_path=out,figsize=(8, 6),categories=data_dict["names"], cbar=True, title='My Three-class CF Matrix')
    with open(str(out / opt.task) + '.txt', 'a') as f:
        f.write("accuracy:{},precision:{},recall:{},f1-scores:{} \n".format(accuracy, precision, recall,
                                                                            f1score))  # label format


if __name__ == '__main__':
    parser = argparse.ArgumentParser(prog='test.py')
    parser.add_argument('--model-name', type=str, default='shufflenet_v2_x0_5', help='initial weights path')  # 要使用那个模型来训练
    parser.add_argument('--weights', nargs='+', type=str, default='runs/exp17_tyre_shufflenet_v2_x0_5/weights/best.pt',
                        help='model.pt path(s)')
    parser.add_argument('--data', type=str, default='data/tyre.yaml', help='data.yaml path')
    parser.add_argument('--batch-size', type=int, default=1, help='size of each image batch')
    parser.add_argument('--img-size', type=int, default=[512, 128], help='inference size (pixels)')
    parser.add_argument('--task', default='test', help="'val', 'test', 'study'")
    parser.add_argument('--device', default='1', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--save-error-img', action='store_true', default=True,
                        help='use weighted image selection for training')  # 是否保存预测错的图片
    opt = parser.parse_args()

    device = select_device(opt.device)

    test(opt, device)
